#ifndef OCK_MEMORY_BRIDGE_MULTI_DEVICE_HETERO_MEMORY_MGR_H
#define OCK_MEMORY_BRIDGE_MULTI_DEVICE_HETERO_MEMORY_MGR_H
#include <memory>
#include <thread>
#include <vector>
#include "ock/hmm/mgr/OckHmmErrorCode.h"
#include "ock/hmm/mgr/OckHmmHMObject.h"
#include "ock/hmm/mgr/OckHmmHeteroMemoryLocation.h"
#include "ock/hmm/mgr/OckHmmMemorySpecification.h"
#include "ock/hmm/mgr/OckHmmTrafficStatisticsInfo.h"
#include "ock/hmm/mgr/OckHmmMemoryPool.h"

namespace ock {
namespace hmm {

struct OckHmmPureDeviceInfo
{
    OckHmmDeviceId deviceId;
    cpu_set_t cpuSet;
    OckHmmMemorySpecification memorySpec;
};

struct OckHmmDeviceInfo
{
    OckHmmDeviceId deviceId{65535U};
    cpu_set_t cpuSet;
    uint32_t transferThreadNum = {0};
    OckHmmMemorySpecification memorySpec;
};

std::ostream &operator<<(std::ostream &os, OckHmmDeviceInfo &data);
using OckHmmDeviceInfoVec = std::vector<OckHmmDeviceInfo>;
std::ostream &operator<<(std::ostream &os, OckHmmDeviceInfoVec &data);
std::ostream &operator<<(std::ostream &os, OckHmmPureDeviceInfo &data);

using OckHmmPureDeviceInfoVec = std::vector<OckHmmPureDeviceInfo>;
std::ostream &operator<<(std::ostream &os, OckHmmPureDeviceInfoVec &data);
std::ostream &operator<<(std::ostream &os, cpu_set_t &cpuSet);

class OckHmmHeteroMemoryMgrBase : public OckHmmMemoryPool
{
public:
    virtual ~OckHmmHeteroMemoryMgrBase() noexcept = default;

    virtual std::pair<OckHmmErrorCode, std::shared_ptr<OckHmmHMObject>> Alloc(
        uint64_t hmoBytes, OckHmmMemoryAllocatePolicy policy = OckHmmMemoryAllocatePolicy::DEVICE_DDR_FIRST
    ) = 0;
    virtual void Free(std::shared_ptr<OckHmmHMObject> hmo) = 0;
    virtual OckHmmErrorCode CopyHMO(
        OckHmmHMObject &dstHMO, uint64_t dstOffset, OckHmmHMObject &srcHMO, uint64_t srcOffset, size_t length) = 0;

    virtual std::shared_ptr<OckHmmResourceUsedInfo> GetUsedInfo(uint64_t fragThresthold) const = 0;
    virtual std::shared_ptr<OckHmmTrafficStatisticsInfo> GetTrafficStatisticsInfo (uint32_t maxGapMilliSeconds = 10) = 0;
};

class OckHmmHeteroMemoryMgr : public OckHmmHeteroMemoryMgrBase
{
public:
    virtual ~OckHmmHeteroMemoryMgr() noexcept = default;
    virtual const OckHmmMemorySpecification &GetSpecific(void) const = 0;
    virtual const cpu_set_t &GetCpuSet(void) const = 0;
    virtual OckHmmDeviceId GetDeviceId(void) const = 0;

    virtual OckHmmErrorCode IncBindMemory(OckHmmHeteroMemoryLocation allocType, uint64_t byteSize,
        uint32_t timeout = 0) = 0;
};

class OckHmmMultiDeviceHeteroMemoryMgrBase
{
public:
    virtual ~OckHmmMultiDeviceHeteroMemoryMgrBase() noexcept = default;
    virtual std::shared_ptr<OckHmmResourceUsedInfo> GetUsedInfo(
        uint64_t fragThresthold, OckHmmDeviceId deviceId) const = 0;
    virtual std::shared_ptr<OckHmmTrafficStatisticsInfo> GetTrafficStatisticsInfo (
        OckHmmDeviceId deviceId, uint32_t maxGapMilliSeconds = 10) = 0;
    virtual const cpu_set_t &GetCpuSet(OckHmmDeviceId deviceId) const = 0;
    virtual std::pair<OckHmmErrorCode, std::shared_ptr<OckHmmHMObject>> Alloc(OckHmmDeviceId deviceId,
        uint64_t hmoBytes, OckHmmMemoryAllocatePolicy policy = OckHmmMemoryAllocatePolicy::DEVICE_DDR_FIRST
    ) = 0;
};

class OckHmmComposeDeviceMgr : public OckHmmMultiDeviceHeteroMemoryMgrBase, public OckHmmHeteroMemoryMgrBase
{
public:
    using OckHmmHeteroMemoryMgrBase::Alloc;
    using OckHmmMultiDeviceHeteroMemoryMgrBase::Alloc;
    using OckHmmHeteroMemoryMgrBase::GetUsedInfo;
    using OckHmmMultiDeviceHeteroMemoryMgrBase::GetUsedInfo;
    using OckHmmHeteroMemoryMgrBase::GetTrafficStatisticsInfo;
    using OckHmmMultiDeviceHeteroMemoryMgrBase::GetTrafficStatisticsInfo;

    virtual ~OckHmmComposeDeviceMgr() noexcept = default;
    virtual const OckHmmMemorySpecification &GetSpecific(OckHmmDeviceId deviceId) const = 0;
};

class OckHmmShareDeviceMgr : public OckHmmHeteroMemoryMgrBase
{
public:
    using OckHmmHeteroMemoryMgrBase::Alloc;
    using OckHmmMultiDeviceHeteroMemoryMgrBase::Alloc;
    using OckHmmHeteroMemoryMgrBase::GetUsedInfo;
    using OckHmmMultiDeviceHeteroMemoryMgrBase::GetUsedInfo;
    using OckHmmHeteroMemoryMgrBase::GetTrafficStatisticsInfo;
    using OckHmmMultiDeviceHeteroMemoryMgrBase::GetTrafficStatisticsInfo;

    virtual ~OckHmmShareDeviceMgr() noexcept = default;
    virtual const OckHmmMemorySpecification &GetDeviceSpecific(OckHmmDeviceId deviceId) const = 0;
    virtual const OckHmmMemorySpecification &GetHostSpecific(void) const = 0;

};

std::ostream &operator<<(std::ostream &os, cpu_set_t &cpuSet);

}
}
#endif